"""
分区策略管理器
提供按时间和股票代码的自动分区管理功能
"""

import asyncio
import time
from typing import List, Dict, Any, Optional, Union, Tuple
from datetime import datetime, date, timedelta
from dataclasses import dataclass, field
from enum import Enum
from collections import defaultdict
from sqlalchemy import select, func, and_, or_, text, MetaData, Table, Column
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.sql import Select
from sqlalchemy.schema import CreateTable, DropTable

from .database import DatabaseManager
from .cache import CacheManager, CacheLevel
from .config import config


class PartitionType(Enum):
    """分区类型"""
    RANGE = "RANGE"           # 范围分区
    HASH = "HASH"             # 哈希分区
    LIST = "LIST"             # 列表分区
    TIME_RANGE = "TIME_RANGE" # 时间范围分区


class PartitionStrategy(Enum):
    """分区策略"""
    DAILY = "DAILY"           # 按日分区
    WEEKLY = "WEEKLY"         # 按周分区
    MONTHLY = "MONTHLY"       # 按月分区
    QUARTERLY = "QUARTERLY"   # 按季度分区
    YEARLY = "YEARLY"         # 按年分区
    BY_STOCK_CODE = "BY_STOCK_CODE"  # 按股票代码分区
    HYBRID = "HYBRID"         # 混合分区策略


class PartitionStatus(Enum):
    """分区状态"""
    ACTIVE = "ACTIVE"         # 活跃
    INACTIVE = "INACTIVE"     # 非活跃
    ARCHIVED = "ARCHIVED"     # 已归档
    DROPPED = "DROPPED"       # 已删除


@dataclass
class PartitionConfig:
    """分区配置"""
    table_name: str
    partition_type: PartitionType
    strategy: PartitionStrategy
    partition_column: str
    retention_days: int = 365  # 数据保留天数
    auto_create: bool = True   # 自动创建分区
    auto_drop: bool = False    # 自动删除过期分区
    partition_size_limit: int = 10000000  # 分区大小限制（行数）
    compression_enabled: bool = True      # 启用压缩
    index_columns: List[str] = field(default_factory=list)
    
    # 时间分区特定配置
    time_format: str = "%Y%m%d"
    advance_partitions: int = 7  # 提前创建分区数量
    
    # 哈希分区特定配置
    hash_partitions: int = 16    # 哈希分区数量
    
    # 股票代码分区特定配置
    stock_codes_per_partition: int = 100


@dataclass
class PartitionInfo:
    """分区信息"""
    partition_name: str
    table_name: str
    partition_type: PartitionType
    partition_key: str
    start_value: Optional[str] = None
    end_value: Optional[str] = None
    row_count: int = 0
    size_mb: float = 0.0
    status: PartitionStatus = PartitionStatus.ACTIVE
    created_at: datetime = field(default_factory=datetime.now)
    last_accessed: Optional[datetime] = None


class TimePartitionManager:
    """时间分区管理器"""
    
    def __init__(self, db_manager: DatabaseManager):
        self.db_manager = db_manager
    
    async def create_time_partitions(
        self,
        session: AsyncSession,
        config: PartitionConfig,
        start_date: date,
        end_date: date
    ) -> List[str]:
        """创建时间分区"""
        created_partitions = []
        current_date = start_date
        
        while current_date <= end_date:
            partition_name = self._generate_time_partition_name(config, current_date)
            
            if await self._partition_exists(session, partition_name):
                current_date = self._get_next_date(current_date, config.strategy)
                continue
            
            # 计算分区范围
            partition_start, partition_end = self._get_partition_range(current_date, config.strategy)
            
            # 创建分区表
            success = await self._create_time_partition_table(
                session, config, partition_name, partition_start, partition_end
            )
            
            if success:
                created_partitions.append(partition_name)
            
            current_date = self._get_next_date(current_date, config.strategy)
        
        return created_partitions
    
    def _generate_time_partition_name(self, config: PartitionConfig, date_value: date) -> str:
        """生成时间分区名称"""
        if config.strategy == PartitionStrategy.DAILY:
            suffix = date_value.strftime("%Y%m%d")
        elif config.strategy == PartitionStrategy.WEEKLY:
            year, week, _ = date_value.isocalendar()
            suffix = f"{year}w{week:02d}"
        elif config.strategy == PartitionStrategy.MONTHLY:
            suffix = date_value.strftime("%Y%m")
        elif config.strategy == PartitionStrategy.QUARTERLY:
            quarter = (date_value.month - 1) // 3 + 1
            suffix = f"{date_value.year}q{quarter}"
        elif config.strategy == PartitionStrategy.YEARLY:
            suffix = date_value.strftime("%Y")
        else:
            suffix = date_value.strftime("%Y%m%d")
        
        return f"{config.table_name}_p_{suffix}"
    
    def _get_partition_range(self, date_value: date, strategy: PartitionStrategy) -> Tuple[date, date]:
        """获取分区范围"""
        if strategy == PartitionStrategy.DAILY:
            start_date = date_value
            end_date = date_value + timedelta(days=1)
        elif strategy == PartitionStrategy.WEEKLY:
            # 获取周一作为开始日期
            days_since_monday = date_value.weekday()
            start_date = date_value - timedelta(days=days_since_monday)
            end_date = start_date + timedelta(days=7)
        elif strategy == PartitionStrategy.MONTHLY:
            start_date = date_value.replace(day=1)
            if start_date.month == 12:
                end_date = start_date.replace(year=start_date.year + 1, month=1)
            else:
                end_date = start_date.replace(month=start_date.month + 1)
        elif strategy == PartitionStrategy.QUARTERLY:
            quarter = (date_value.month - 1) // 3 + 1
            start_month = (quarter - 1) * 3 + 1
            start_date = date_value.replace(month=start_month, day=1)
            if quarter == 4:
                end_date = start_date.replace(year=start_date.year + 1, month=1)
            else:
                end_date = start_date.replace(month=start_month + 3)
        elif strategy == PartitionStrategy.YEARLY:
            start_date = date_value.replace(month=1, day=1)
            end_date = start_date.replace(year=start_date.year + 1)
        else:
            start_date = date_value
            end_date = date_value + timedelta(days=1)
        
        return start_date, end_date
    
    def _get_next_date(self, current_date: date, strategy: PartitionStrategy) -> date:
        """获取下一个日期"""
        if strategy == PartitionStrategy.DAILY:
            return current_date + timedelta(days=1)
        elif strategy == PartitionStrategy.WEEKLY:
            return current_date + timedelta(days=7)
        elif strategy == PartitionStrategy.MONTHLY:
            if current_date.month == 12:
                return current_date.replace(year=current_date.year + 1, month=1)
            else:
                return current_date.replace(month=current_date.month + 1)
        elif strategy == PartitionStrategy.QUARTERLY:
            quarter = (current_date.month - 1) // 3 + 1
            if quarter == 4:
                return current_date.replace(year=current_date.year + 1, month=1)
            else:
                next_month = quarter * 3 + 1
                return current_date.replace(month=next_month)
        elif strategy == PartitionStrategy.YEARLY:
            return current_date.replace(year=current_date.year + 1)
        else:
            return current_date + timedelta(days=1)
    
    async def _create_time_partition_table(
        self,
        session: AsyncSession,
        config: PartitionConfig,
        partition_name: str,
        start_date: date,
        end_date: date
    ) -> bool:
        """创建时间分区表"""
        try:
            # 构建分区表创建SQL
            create_sql = f"""
                CREATE TABLE {partition_name} (
                    LIKE {config.table_name} INCLUDING ALL
                ) INHERITS ({config.table_name});
            """
            
            await session.execute(text(create_sql))
            
            # 添加分区约束
            constraint_sql = f"""
                ALTER TABLE {partition_name} 
                ADD CONSTRAINT {partition_name}_check 
                CHECK ({config.partition_column} >= '{start_date}' 
                       AND {config.partition_column} < '{end_date}');
            """
            
            await session.execute(text(constraint_sql))
            
            # 创建索引
            if config.index_columns:
                for i, col in enumerate(config.index_columns):
                    index_sql = f"""
                        CREATE INDEX idx_{partition_name}_{col} 
                        ON {partition_name} ({col});
                    """
                    await session.execute(text(index_sql))
            
            # 创建分区列索引
            partition_index_sql = f"""
                CREATE INDEX idx_{partition_name}_{config.partition_column} 
                ON {partition_name} ({config.partition_column});
            """
            await session.execute(text(partition_index_sql))
            
            await session.commit()
            return True
            
        except Exception as e:
            await session.rollback()
            print(f"创建时间分区表失败: {e}")
            return False
    
    async def _partition_exists(self, session: AsyncSession, partition_name: str) -> bool:
        """检查分区是否存在"""
        check_sql = """
            SELECT EXISTS (
                SELECT 1 FROM information_schema.tables 
                WHERE table_name = :partition_name
            );
        """
        result = await session.execute(text(check_sql), {"partition_name": partition_name})
        scalar_result = result.scalar()
        return scalar_result is not None and bool(scalar_result)


class HashPartitionManager:
    """哈希分区管理器"""
    
    def __init__(self, db_manager: DatabaseManager):
        self.db_manager = db_manager
    
    async def create_hash_partitions(
        self,
        session: AsyncSession,
        config: PartitionConfig
    ) -> List[str]:
        """创建哈希分区"""
        created_partitions = []
        
        for i in range(config.hash_partitions):
            partition_name = f"{config.table_name}_p_hash_{i:02d}"
            
            if await self._partition_exists(session, partition_name):
                continue
            
            success = await self._create_hash_partition_table(
                session, config, partition_name, i
            )
            
            if success:
                created_partitions.append(partition_name)
        
        return created_partitions
    
    async def _create_hash_partition_table(
        self,
        session: AsyncSession,
        config: PartitionConfig,
        partition_name: str,
        hash_value: int
    ) -> bool:
        """创建哈希分区表"""
        try:
            # 构建分区表创建SQL
            create_sql = f"""
                CREATE TABLE {partition_name} (
                    LIKE {config.table_name} INCLUDING ALL
                ) INHERITS ({config.table_name});
            """
            
            await session.execute(text(create_sql))
            
            # 添加哈希约束
            constraint_sql = f"""
                ALTER TABLE {partition_name} 
                ADD CONSTRAINT {partition_name}_check 
                CHECK (hashtext({config.partition_column}) % {config.hash_partitions} = {hash_value});
            """
            
            await session.execute(text(constraint_sql))
            
            # 创建索引
            if config.index_columns:
                for col in config.index_columns:
                    index_sql = f"""
                        CREATE INDEX idx_{partition_name}_{col} 
                        ON {partition_name} ({col});
                    """
                    await session.execute(text(index_sql))
            
            await session.commit()
            return True
            
        except Exception as e:
            await session.rollback()
            print(f"创建哈希分区表失败: {e}")
            return False
    
    async def _partition_exists(self, session: AsyncSession, partition_name: str) -> bool:
        """检查分区是否存在"""
        check_sql = """
            SELECT EXISTS (
                SELECT 1 FROM information_schema.tables 
                WHERE table_name = :partition_name
            );
        """
        result = await session.execute(text(check_sql), {"partition_name": partition_name})
        scalar_result = result.scalar()
        return scalar_result is not None and bool(scalar_result)


class StockCodePartitionManager:
    """股票代码分区管理器"""
    
    def __init__(self, db_manager: DatabaseManager):
        self.db_manager = db_manager
    
    async def create_stock_code_partitions(
        self,
        session: AsyncSession,
        config: PartitionConfig,
        stock_codes: List[str]
    ) -> List[str]:
        """创建股票代码分区"""
        created_partitions = []
        
        # 按股票代码数量分组
        code_groups = self._group_stock_codes(stock_codes, config.stock_codes_per_partition)
        
        for i, code_group in enumerate(code_groups):
            partition_name = f"{config.table_name}_p_stocks_{i:03d}"
            
            if await self._partition_exists(session, partition_name):
                continue
            
            success = await self._create_stock_partition_table(
                session, config, partition_name, code_group
            )
            
            if success:
                created_partitions.append(partition_name)
        
        return created_partitions
    
    def _group_stock_codes(self, stock_codes: List[str], group_size: int) -> List[List[str]]:
        """将股票代码分组"""
        groups = []
        for i in range(0, len(stock_codes), group_size):
            groups.append(stock_codes[i:i + group_size])
        return groups
    
    async def _create_stock_partition_table(
        self,
        session: AsyncSession,
        config: PartitionConfig,
        partition_name: str,
        stock_codes: List[str]
    ) -> bool:
        """创建股票代码分区表"""
        try:
            # 构建分区表创建SQL
            create_sql = f"""
                CREATE TABLE {partition_name} (
                    LIKE {config.table_name} INCLUDING ALL
                ) INHERITS ({config.table_name});
            """
            
            await session.execute(text(create_sql))
            
            # 添加股票代码约束
            codes_list = "', '".join(stock_codes)
            constraint_sql = f"""
                ALTER TABLE {partition_name} 
                ADD CONSTRAINT {partition_name}_check 
                CHECK ({config.partition_column} IN ('{codes_list}'));
            """
            
            await session.execute(text(constraint_sql))
            
            # 创建索引
            if config.index_columns:
                for col in config.index_columns:
                    index_sql = f"""
                        CREATE INDEX idx_{partition_name}_{col} 
                        ON {partition_name} ({col});
                    """
                    await session.execute(text(index_sql))
            
            # 创建股票代码索引
            stock_index_sql = f"""
                CREATE INDEX idx_{partition_name}_{config.partition_column} 
                ON {partition_name} ({config.partition_column});
            """
            await session.execute(text(stock_index_sql))
            
            await session.commit()
            return True
            
        except Exception as e:
            await session.rollback()
            print(f"创建股票代码分区表失败: {e}")
            return False
    
    async def _partition_exists(self, session: AsyncSession, partition_name: str) -> bool:
        """检查分区是否存在"""
        check_sql = """
            SELECT EXISTS (
                SELECT 1 FROM information_schema.tables 
                WHERE table_name = :partition_name
            );
        """
        result = await session.execute(text(check_sql), {"partition_name": partition_name})
        scalar_result = result.scalar()
        return scalar_result is not None and bool(scalar_result)


class PartitionManager:
    """分区管理器主类"""
    
    def __init__(self, db_manager: DatabaseManager, cache_manager: CacheManager):
        self.db_manager = db_manager
        self.cache_manager = cache_manager
        
        self.time_manager = TimePartitionManager(db_manager)
        self.hash_manager = HashPartitionManager(db_manager)
        self.stock_manager = StockCodePartitionManager(db_manager)
        
        self.partition_configs = {}
        self.partition_info = defaultdict(list)
        self.maintenance_tasks = {}
    
    def register_partition_config(self, config: PartitionConfig):
        """注册分区配置"""
        self.partition_configs[config.table_name] = config
    
    async def create_partitions(
        self,
        table_name: str,
        start_date: Optional[date] = None,
        end_date: Optional[date] = None,
        stock_codes: Optional[List[str]] = None
    ) -> List[str]:
        """创建分区"""
        if table_name not in self.partition_configs:
            raise ValueError(f"表 {table_name} 的分区配置不存在")
        
        config = self.partition_configs[table_name]
        created_partitions = []
        
        async with self.db_manager.get_session() as session:
            if config.partition_type == PartitionType.TIME_RANGE:
                if not start_date:
                    start_date = date.today()
                if not end_date:
                    end_date = start_date + timedelta(days=config.advance_partitions)
                
                partitions = await self.time_manager.create_time_partitions(
                    session, config, start_date, end_date
                )
                created_partitions.extend(partitions)
            
            elif config.partition_type == PartitionType.HASH:
                partitions = await self.hash_manager.create_hash_partitions(session, config)
                created_partitions.extend(partitions)
            
            elif config.strategy == PartitionStrategy.BY_STOCK_CODE:
                if not stock_codes:
                    # 获取所有活跃股票代码
                    stock_codes = await self._get_active_stock_codes(session)
                
                partitions = await self.stock_manager.create_stock_code_partitions(
                    session, config, stock_codes
                )
                created_partitions.extend(partitions)
        
        # 更新分区信息缓存
        await self._update_partition_info_cache(table_name, created_partitions)
        
        return created_partitions
    
    async def _get_active_stock_codes(self, session: AsyncSession) -> List[str]:
        """获取活跃股票代码"""
        # 这里应该根据实际的股票信息表查询
        # 简化实现，返回示例代码
        sql = """
            SELECT DISTINCT stock_code 
            FROM stock_basic_info 
            WHERE is_active = true 
            ORDER BY stock_code
        """
        try:
            result = await session.execute(text(sql))
            return [row[0] for row in result.fetchall()]
        except:
            # 如果表不存在，返回示例代码
            return [f"{i:06d}" for i in range(1, 5001)]  # 000001-005000
    
    async def drop_expired_partitions(self, table_name: str) -> List[str]:
        """删除过期分区"""
        if table_name not in self.partition_configs:
            return []
        
        config = self.partition_configs[table_name]
        if not config.auto_drop:
            return []
        
        dropped_partitions = []
        cutoff_date = date.today() - timedelta(days=config.retention_days)
        
        async with self.db_manager.get_session() as session:
            # 获取所有分区
            partitions = await self._get_table_partitions(session, table_name)
            
            for partition in partitions:
                if self._is_partition_expired(partition, cutoff_date, config):
                    success = await self._drop_partition(session, partition)
                    if success:
                        dropped_partitions.append(partition)
        
        return dropped_partitions
    
    async def _get_table_partitions(self, session: AsyncSession, table_name: str) -> List[str]:
        """获取表的所有分区"""
        sql = """
            SELECT schemaname, tablename 
            FROM pg_tables 
            WHERE tablename LIKE :pattern
        """
        result = await session.execute(text(sql), {"pattern": f"{table_name}_p_%"})
        return [row.tablename for row in result.fetchall()]
    
    def _is_partition_expired(
        self, 
        partition_name: str, 
        cutoff_date: date, 
        config: PartitionConfig
    ) -> bool:
        """检查分区是否过期"""
        # 从分区名称解析日期
        try:
            if "_p_" in partition_name:
                date_part = partition_name.split("_p_")[1]
                
                if config.strategy == PartitionStrategy.DAILY:
                    partition_date = datetime.strptime(date_part, "%Y%m%d").date()
                elif config.strategy == PartitionStrategy.MONTHLY:
                    partition_date = datetime.strptime(date_part, "%Y%m").date()
                elif config.strategy == PartitionStrategy.YEARLY:
                    partition_date = datetime.strptime(date_part, "%Y").date()
                else:
                    return False
                
                return partition_date < cutoff_date
        except:
            pass
        
        return False
    
    async def _drop_partition(self, session: AsyncSession, partition_name: str) -> bool:
        """删除分区"""
        try:
            drop_sql = f"DROP TABLE IF EXISTS {partition_name};"
            await session.execute(text(drop_sql))
            await session.commit()
            return True
        except Exception as e:
            await session.rollback()
            print(f"删除分区 {partition_name} 失败: {e}")
            return False
    
    async def get_partition_statistics(self, table_name: str) -> Dict[str, Any]:
        """获取分区统计信息"""
        async with self.db_manager.get_session() as session:
            partitions = await self._get_table_partitions(session, table_name)
            
            stats = {
                'total_partitions': len(partitions),
                'partition_details': []
            }
            
            for partition in partitions:
                partition_stats = await self._get_partition_stats(session, partition)
                stats['partition_details'].append(partition_stats)
            
            return stats
    
    async def _get_partition_stats(self, session: AsyncSession, partition_name: str) -> Dict[str, Any]:
        """获取单个分区统计信息"""
        try:
            # 获取行数
            count_sql = f"SELECT COUNT(*) FROM {partition_name};"
            count_result = await session.execute(text(count_sql))
            row_count = count_result.scalar()
            
            # 获取表大小
            size_sql = f"SELECT pg_total_relation_size('{partition_name}');"
            size_result = await session.execute(text(size_sql))
            size_bytes = size_result.scalar() or 0
            size_mb = size_bytes / (1024 * 1024)
            
            return {
                'partition_name': partition_name,
                'row_count': row_count,
                'size_mb': round(size_mb, 2),
                'last_updated': datetime.now()
            }
        except Exception as e:
            return {
                'partition_name': partition_name,
                'row_count': 0,
                'size_mb': 0.0,
                'error': str(e)
            }
    
    async def _update_partition_info_cache(self, table_name: str, partitions: List[str]):
        """更新分区信息缓存"""
        cache_key = f"partitions:{table_name}"
        partition_data = {
            'partitions': partitions,
            'updated_at': datetime.now().isoformat()
        }
        
        cache = self.cache_manager.get_cache()
        await cache.set(
            cache_key,
            str(partition_data),
            ttl=3600  # 1小时缓存
        )
    
    async def setup_maintenance_tasks(self):
        """设置维护任务"""
        for table_name, config in self.partition_configs.items():
            if config.auto_create:
                # 创建自动创建分区的任务
                task = asyncio.create_task(self._auto_create_partitions_task(table_name))
                self.maintenance_tasks[f"auto_create_{table_name}"] = task
            
            if config.auto_drop:
                # 创建自动删除过期分区的任务
                task = asyncio.create_task(self._auto_drop_partitions_task(table_name))
                self.maintenance_tasks[f"auto_drop_{table_name}"] = task
    
    async def _auto_create_partitions_task(self, table_name: str):
        """自动创建分区任务"""
        while True:
            try:
                config = self.partition_configs[table_name]
                
                # 检查是否需要创建新分区
                future_date = date.today() + timedelta(days=config.advance_partitions)
                await self.create_partitions(table_name, end_date=future_date)
                
                # 每天检查一次
                await asyncio.sleep(86400)
                
            except Exception as e:
                print(f"自动创建分区任务失败: {e}")
                await asyncio.sleep(3600)  # 出错后1小时重试
    
    async def _auto_drop_partitions_task(self, table_name: str):
        """自动删除过期分区任务"""
        while True:
            try:
                await self.drop_expired_partitions(table_name)
                
                # 每天检查一次
                await asyncio.sleep(86400)
                
            except Exception as e:
                print(f"自动删除分区任务失败: {e}")
                await asyncio.sleep(3600)  # 出错后1小时重试
    
    def stop_maintenance_tasks(self):
        """停止维护任务"""
        for task_name, task in self.maintenance_tasks.items():
            task.cancel()
        self.maintenance_tasks.clear()


# 全局分区管理器实例
_partition_manager: Optional[PartitionManager] = None

def get_partition_manager() -> PartitionManager:
    """获取全局分区管理器实例"""
    global _partition_manager
    if _partition_manager is None:
        from .database import db_manager
        from .cache import cache_manager
        
        _partition_manager = PartitionManager(db_manager, cache_manager)
    
    return _partition_manager


# 预设分区配置
def get_preset_partition_configs() -> List[PartitionConfig]:
    """获取预设分区配置"""
    return [
        # 股票日线数据按时间分区
        PartitionConfig(
            table_name="stock_daily_data",
            partition_type=PartitionType.TIME_RANGE,
            strategy=PartitionStrategy.MONTHLY,
            partition_column="trade_date",
            retention_days=1095,  # 3年
            auto_create=True,
            auto_drop=True,
            advance_partitions=30,  # 提前30天创建
            index_columns=["stock_code", "trade_date"]
        ),
        
        # 分钟级数据按时间分区
        PartitionConfig(
            table_name="stock_minute_data",
            partition_type=PartitionType.TIME_RANGE,
            strategy=PartitionStrategy.DAILY,
            partition_column="trade_datetime",
            retention_days=90,  # 90天
            auto_create=True,
            auto_drop=True,
            advance_partitions=7,
            index_columns=["stock_code", "trade_datetime"]
        ),
        
        # 财务数据按股票代码分区
        PartitionConfig(
            table_name="financial_statements",
            partition_type=PartitionType.LIST,
            strategy=PartitionStrategy.BY_STOCK_CODE,
            partition_column="stock_code",
            retention_days=3650,  # 10年
            auto_create=True,
            auto_drop=False,
            stock_codes_per_partition=200,
            index_columns=["stock_code", "report_date", "report_type"]
        ),
        
        # 交易明细按哈希分区
        PartitionConfig(
            table_name="trade_details",
            partition_type=PartitionType.HASH,
            strategy=PartitionStrategy.HYBRID,
            partition_column="trade_id",
            retention_days=365,
            auto_create=True,
            auto_drop=True,
            hash_partitions=16,
            index_columns=["trade_id", "stock_code", "trade_time"]
        ),
        
        # 资金流向数据按时间和股票代码混合分区
        PartitionConfig(
            table_name="money_flow_data",
            partition_type=PartitionType.TIME_RANGE,
            strategy=PartitionStrategy.WEEKLY,
            partition_column="trade_date",
            retention_days=730,  # 2年
            auto_create=True,
            auto_drop=True,
            advance_partitions=14,  # 提前2周创建
            index_columns=["stock_code", "trade_date", "flow_type"]
        )
    ]


class BatchPartitionOperator:
    """批量分区操作器"""
    
    def __init__(self, partition_manager: PartitionManager):
        self.partition_manager = partition_manager
    
    async def batch_create_partitions(
        self,
        table_configs: List[Tuple[str, date, date]]
    ) -> Dict[str, List[str]]:
        """批量创建分区"""
        results = {}
        
        for table_name, start_date, end_date in table_configs:
            try:
                partitions = await self.partition_manager.create_partitions(
                    table_name, start_date, end_date
                )
                results[table_name] = partitions
            except Exception as e:
                results[table_name] = []
                print(f"批量创建分区失败 {table_name}: {e}")
        
        return results
    
    async def batch_drop_partitions(
        self,
        partition_names: List[str]
    ) -> Dict[str, bool]:
        """批量删除分区"""
        results = {}
        
        async with self.partition_manager.db_manager.get_session() as session:
            for partition_name in partition_names:
                try:
                    success = await self.partition_manager._drop_partition(session, partition_name)
                    results[partition_name] = success
                except Exception as e:
                    results[partition_name] = False
                    print(f"批量删除分区失败 {partition_name}: {e}")
        
        return results
    
    async def batch_get_partition_statistics(
        self,
        table_names: List[str]
    ) -> Dict[str, Dict[str, Any]]:
        """批量获取分区统计信息"""
        results = {}
        
        for table_name in table_names:
            try:
                stats = await self.partition_manager.get_partition_statistics(table_name)
                results[table_name] = stats
            except Exception as e:
                results[table_name] = {'error': str(e)}
                print(f"批量获取分区统计失败 {table_name}: {e}")
        
        return results


class PartitionHealthChecker:
    """分区健康检查器"""
    
    def __init__(self, partition_manager: PartitionManager):
        self.partition_manager = partition_manager
        self.health_metrics = {}
    
    async def check_partition_health(self, table_name: str) -> Dict[str, Any]:
        """检查分区健康状态"""
        health_report = {
            'table_name': table_name,
            'check_time': datetime.now(),
            'status': 'healthy',
            'issues': [],
            'recommendations': [],
            'metrics': {}
        }
        
        try:
            # 获取分区统计信息
            stats = await self.partition_manager.get_partition_statistics(table_name)
            health_report['metrics'] = stats
            
            # 检查分区数量
            total_partitions = stats.get('total_partitions', 0)
            if total_partitions == 0:
                health_report['issues'].append('没有找到分区')
                health_report['status'] = 'warning'
                health_report['recommendations'].append('创建初始分区')
            
            # 检查分区大小分布
            partition_details = stats.get('partition_details', [])
            if partition_details:
                sizes = [p.get('size_mb', 0) for p in partition_details]
                avg_size = sum(sizes) / len(sizes) if sizes else 0
                max_size = max(sizes) if sizes else 0
                
                # 检查是否有过大的分区
                config = self.partition_manager.partition_configs.get(table_name)
                if config and max_size > 1000:  # 1GB
                    health_report['issues'].append(f'存在过大分区: {max_size:.2f}MB')
                    health_report['status'] = 'warning'
                    health_report['recommendations'].append('考虑调整分区策略')
                
                # 检查空分区
                empty_partitions = [p for p in partition_details if p.get('row_count', 0) == 0]
                if len(empty_partitions) > 5:
                    health_report['issues'].append(f'存在{len(empty_partitions)}个空分区')
                    health_report['recommendations'].append('清理空分区')
            
            # 检查分区配置
            if table_name in self.partition_manager.partition_configs:
                config = self.partition_manager.partition_configs[table_name]
                
                # 检查自动维护任务状态
                if config.auto_create:
                    task_name = f"auto_create_{table_name}"
                    if task_name not in self.partition_manager.maintenance_tasks:
                        health_report['issues'].append('自动创建分区任务未运行')
                        health_report['recommendations'].append('启动自动维护任务')
                
                if config.auto_drop:
                    task_name = f"auto_drop_{table_name}"
                    if task_name not in self.partition_manager.maintenance_tasks:
                        health_report['issues'].append('自动删除分区任务未运行')
                        health_report['recommendations'].append('启动自动维护任务')
            
            # 设置最终状态
            if health_report['issues']:
                if health_report['status'] != 'warning':
                    health_report['status'] = 'needs_attention'
            
        except Exception as e:
            health_report['status'] = 'error'
            health_report['issues'].append(f'健康检查失败: {str(e)}')
        
        return health_report
    
    async def check_all_partitions_health(self) -> Dict[str, Dict[str, Any]]:
        """检查所有分区的健康状态"""
        health_reports = {}
        
        for table_name in self.partition_manager.partition_configs.keys():
            health_reports[table_name] = await self.check_partition_health(table_name)
        
        return health_reports
    
    async def get_health_summary(self) -> Dict[str, Any]:
        """获取健康状态摘要"""
        all_reports = await self.check_all_partitions_health()
        
        summary = {
            'total_tables': len(all_reports),
            'healthy_tables': 0,
            'warning_tables': 0,
            'error_tables': 0,
            'total_issues': 0,
            'common_issues': defaultdict(int),
            'check_time': datetime.now()
        }
        
        for table_name, report in all_reports.items():
            status = report.get('status', 'unknown')
            if status == 'healthy':
                summary['healthy_tables'] += 1
            elif status == 'warning':
                summary['warning_tables'] += 1
            elif status in ['error', 'needs_attention']:
                summary['error_tables'] += 1
            
            issues = report.get('issues', [])
            summary['total_issues'] += len(issues)
            
            for issue in issues:
                summary['common_issues'][issue] += 1
        
        return summary


# 辅助函数
async def initialize_preset_partitions():
    """初始化预设分区配置"""
    partition_manager = get_partition_manager()
    
    # 注册预设配置
    for config in get_preset_partition_configs():
        partition_manager.register_partition_config(config)
    
    # 启动维护任务
    await partition_manager.setup_maintenance_tasks()
    
    print("预设分区配置已初始化")


async def create_initial_partitions():
    """创建初始分区"""
    partition_manager = get_partition_manager()
    
    # 为每个配置的表创建初始分区
    for table_name in partition_manager.partition_configs.keys():
        try:
            partitions = await partition_manager.create_partitions(table_name)
            print(f"为表 {table_name} 创建了 {len(partitions)} 个分区")
        except Exception as e:
            print(f"为表 {table_name} 创建分区失败: {e}")


async def cleanup_old_partitions():
    """清理旧分区"""
    partition_manager = get_partition_manager()
    
    for table_name in partition_manager.partition_configs.keys():
        try:
            dropped = await partition_manager.drop_expired_partitions(table_name)
            if dropped:
                print(f"为表 {table_name} 删除了 {len(dropped)} 个过期分区")
        except Exception as e:
            print(f"清理表 {table_name} 的过期分区失败: {e}")


def get_batch_operator() -> BatchPartitionOperator:
    """获取批量分区操作器"""
    partition_manager = get_partition_manager()
    return BatchPartitionOperator(partition_manager)


def get_health_checker() -> PartitionHealthChecker:
    """获取分区健康检查器"""
    partition_manager = get_partition_manager()
    return PartitionHealthChecker(partition_manager)


# 导出全局实例供测试使用
partition_manager = get_partition_manager()